import numpy as np
from sklearn.decomposition import NMF
import matplotlib.pyplot as plt 
import numpy as np
import h5py
from mpl_toolkits.axes_grid1 import make_axes_locatable
  
def read_mat_hyperspectral_image(file_path, dataset_name):  
    try:  
        with h5py.File(file_path, 'r') as f:  
            dataset = f[dataset_name]  
            data = np.array(dataset)  
            return data  
    except Exception as e:  
        print(f"An error occurred while reading the dataset from {file_path}: {e}")  
        return None  
  
# 使用函数读取文件  
file_path = './data/0031_9.mat'  # 替换为你的.mat文件路径  
dataset_name = 'cube'  # 替换为MATLAB中保存图像数据的变量名  
hsi_data = read_mat_hyperspectral_image(file_path, dataset_name)  
  
if hsi_data is None:  
    print("Failed to read the hyperspectral image.")  
else:  
    # 现在high_spectral_image是一个三维NumPy数组，可以进一步处理  
    print("hsi_data.shape:"+str(hsi_data.shape))
channel = hsi_data.shape[0]
height = hsi_data.shape[1]
width = hsi_data.shape[2]
print("channel:" + str(channel) + ",height:" + str(height) + ",width:" + str(width))

# 将数据维度进行变换
hsi_data = np.transpose(hsi_data,(2,1,0))
print("翻转维度的hsi_data的形状为:" + str(hsi_data.shape))

# 重塑数据以适应NMF输入要求
hsi_data_reshaped = hsi_data.reshape(-1,hsi_data.shape[-1])
print("hsi_data_reshaped.shape:" + str(hsi_data_reshaped.shape))

# NMF分解
n_components = 3  # 输出的波段数，即最终的RGB通道数
nmf_model = NMF(n_components=n_components, init='random', random_state=0,max_iter=7000)
W = nmf_model.fit_transform(hsi_data_reshaped)
H = nmf_model.components_

# 重构图像，使用NMF的结果
hsi_nmf = np.dot(W, H).reshape(hsi_data.shape)

# 选择用于RGB合成的波段
red_band = hsi_nmf[:, :, 29]
green_band = hsi_nmf[:, :, 12]
blue_band = hsi_nmf[:, :, 7]

# 线性拉伸
min_val = np.min(hsi_nmf)
max_val = np.max(hsi_nmf)
red_band = (red_band - min_val) / (max_val - min_val)
green_band = (green_band - min_val) / (max_val - min_val)
blue_band = (blue_band - min_val) / (max_val - min_val)

# 合并波段
rgb_image = np.dstack((red_band, green_band, blue_band))

# 显示RGB图像
dpi = 100
fig_width_inches = width / dpi  
fig_height_inches = height / dpi 
fig, ax = plt.subplots(figsize=(fig_height_inches,fig_width_inches))
divider = make_axes_locatable(ax)
ax.imshow(rgb_image)
ax.axis('off') # 关闭坐标轴
print("rgb_image.shape:" + str(rgb_image.shape))
plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
plt.margins(0, 0)
plt.savefig("0031_10.jpg",dpi = dpi)
plt.show()